from typing import Union

from scripts.Log import Log
from scripts.actions.Action import Action
from scripts.device.AdbDevice import AdbDevice
from scripts.device.AdbShell import AdbShell
from scripts.device.FastbootDevice import FastbootDevice


class RebootHelper(object):
    def __init__(self, action: Action):
        self._action = action

    def safe_reboot(self, device: AdbDevice, shell: AdbShell) -> (bool, AdbDevice, AdbShell):
        while True:
            result, code = shell.safe_reboot()
            if result:
                return True, device, shell
            elif code == 137:
                Log.w('safe reboot failed, maybe root failed. try root again.')
                # restart adb server
                device.kill_server()
                device = AdbDevice.wait_for_device(self._action.sn())
                Log.d('got device: %s after adb restarted' % device.sn())

                # root device again
                if not self.root(device, self._action.android_version()):
                    Log.e('root failed')
                    return False, device, shell
                device = self.wait_for_device_rooted()
                shell = device.shell()

    def root(self, device: AdbDevice, android_version: int) -> bool:
        if android_version == 9:
            root_success = self.root9(device)
        elif android_version == 11:
            root_success = self.root11(device)
        else:
            Log.e('unable to root device with android version: %d' % android_version)
            root_success = False
        return root_success

    def root9(self, adb_device: AdbDevice) -> bool:
        # reboot to fastboot
        adb_device.reboot_bootloader()

        # get fastboot device
        fastboot_device = self.wait_for_fastboot_device()

        # do device root
        result = fastboot_device.enable_root_9()

        # reboot
        fastboot_device.reboot()

        return result

    def root11(self, adb_device: AdbDevice) -> bool:
        # reboot to fastboot
        adb_device.reboot_bootloader()

        # get fastboot device
        fastboot_device = self.wait_for_fastboot_device()

        # do device root
        result = fastboot_device.enable_root_11()

        # reboot
        fastboot_device.reboot()

        return result

    def wait_for_fastboot_device(self) -> FastbootDevice:
        fastboot_device = FastbootDevice.wait_for_device(self._action.sn())
        Log.d('got fastboot device: %s' % fastboot_device.sn())
        return fastboot_device

    def wait_for_device_rooted(self) -> Union[None, AdbDevice]:
        device = AdbDevice.wait_for_device(self._action.sn())
        device.adb_root()
        device = AdbDevice.wait_for_device(self._action.sn())
        with device.shell() as shell:
            rooted = shell.is_rooted()
        if rooted:
            Log.d('got rooted device: %s' % device.sn())
            return device
        else:
            Log.d('root device failed: %s' % device.sn())
            return None
